package co.yixiang.modules.system.utils;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.apache.commons.lang.StringUtils;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

public class ChatCompletionsExample {
    //    static String apiKey = System.getenv("");
    static String apiKey = "cb08102a-c849-4cc4-9dab-f08b15cef8f3";
    static String model = "ep-20250227100314-jr2bt";
    static String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";

    // The output time of the reasoning model is relatively long. Please increase the timeout period.
    public static ArkService question(String sessionId, String question) {
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service = ArkService.builder().timeout(Duration.ofSeconds(1800)).connectTimeout(Duration.ofSeconds(20)).dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();
        try {
            System.out.println("\n[Recommended]----- streaming request -----");
            final List<ChatMessage> streamMessages = new ArrayList<>();
            final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(question).build();
            streamMessages.add(streamUserMessage);
            ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                    .model(model)
                    .messages(streamMessages)
                    .build();
            service.streamChatCompletion(streamChatCompletionRequest)
                    .doOnError(Throwable::printStackTrace)
                    .blockingForEach(
                            delta -> {
                                if (!delta.getChoices().isEmpty()) {
                                    if (StringUtils.isNotEmpty(delta.getChoices().get(0).getMessage().getReasoningContent())) {
                                                                            System.out.print(delta.getChoices().get(0).getMessage().getReasoningContent());
                                        WebSocket.sendMessage(delta.getChoices().get(0).getMessage().getReasoningContent(), sessionId);
                                    } else {
                                                                            System.out.print(delta.getChoices().get(0).getMessage().getContent());
                                        WebSocket.sendMessage(delta.getChoices().get(0).getMessage().getContent(), sessionId);
                                    }
                                }
                            }
                    );
            System.out.println("\n----- standard request -----");
            final List<ChatMessage> messages = new ArrayList<>();
            final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(question).build();
            messages.add(userMessage);
            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(model)
                    .messages(messages)
                    .build();
            service.createChatCompletion(chatCompletionRequest).getChoices().forEach(
                    choice -> {
                        //                    System.out.println(choice.getMessage().getReasoningContent());
                        WebSocket.sendMessage(choice.getMessage().getReasoningContent(), sessionId);
                        //                    System.out.println(choice.getMessage().getContent());
                        WebSocket.sendMessage(choice.getMessage().getContent(), sessionId);
                    }
            );
            // shutdown service after all requests is finished
        } catch (Exception e) {
            service.shutdownExecutor();
            throw new RuntimeException(e);
        }
        return service;
    }

    public static String questionAll( String question) {
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service = ArkService.builder().timeout(Duration.ofSeconds(1800)).connectTimeout(Duration.ofSeconds(20)).dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();
        StringBuilder stringBuilder = new StringBuilder();
        try {
            System.out.println("\n----- standard request -----");
            final List<ChatMessage> messages = new ArrayList<>();
//            final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build();
            final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(question).build();
//            messages.add(systemMessage);
            messages.add(userMessage);

            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(model)
                    .messages(messages)
                    .build();
//            long t1 = System.currentTimeMillis();
            service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> stringBuilder.append(choice.getMessage().getContent()));
//            long t2 = System.currentTimeMillis();
            System.out.println(stringBuilder.toString());
//            System.out.println(t2 - t1);
            // shutdown service
            service.shutdownExecutor();
        } catch (Exception e) {
            service.shutdownExecutor();
            throw new RuntimeException(e);
        }
        return stringBuilder.toString();
    }

    public static void main(String[] args) {
        ArkService question = question("1","【郑州四中】总分满分是960分唐文熙初二上学期末成绩:总分704.5，文化课673(总分满分是960分，语文数学英语满分各120分，" +
                "其他科目满分各100分)，语文71，数学111，英语105，物理89.生物93，道德与法治53，历史77地理74，体育分数31.5，" +
                "请耐心帮孩子总结经验教训，以及制定出初二下学期的学习方案");
        question.shutdownExecutor();
    }
}